from torch.utils.data import Dataset  # 导入 PyTorch 的 Dataset 类，用于自定义数据集
from datasets import load_from_disk  # 导入 Hugging Face 的 `load_from_disk` 函数，用于加载缓存数据


# 定义一个自定义数据集类 MyDataset，继承自 PyTorch 的 Dataset 类
class MyDataset(Dataset):
    def __init__(self, split):
        """
        初始化函数，加载数据并根据 split 参数选择数据分割
        :param split: 数据分割类型（"train", "test", 或 "validation"）
        """
        # 从磁盘加载缓存数据集
        self.dataset = load_from_disk(r"D:\SentimentBERT\data\ChnSentiCorp")

        # 根据 split 参数选择对应的数据分割
        if split == 'train':
            self.dataset = self.dataset["train"]  # 加载训练集
        elif split == "test":
            self.dataset = self.dataset["test"]  # 加载测试集
        elif split == "validation":
            self.dataset = self.dataset["validation"]  # 加载验证集

    def __len__(self):
        """
        返回数据集的长度
        """
        return len(self.dataset)  # 返回数据集中样本的数量

    def __getitem__(self, item):
        """
        根据索引获取数据集中的单个样本
        :param item: 样本的索引
        :return: (text, label) 元组，包含文本和标签
        """
        text = self.dataset[item]['text']  # 获取当前索引对应的文本
        label = self.dataset[item]['label']  # 获取当前索引对应的标签

        return text, label  # 返回文本和标签


# 主程序部分
if __name__ == '__main__':
    dataset = MyDataset("test")  # 创建 MyDataset 实例，加载测试集
    for data in dataset:  # 遍历数据集中的每个样本
        print(data)  # 打印每个样本的内容（文本和标签）
   # 逻辑的实现
   # 1.自定义数据集类
       # 通过继承 PyTorch 的 Dataset 类，实现了一个自定义数据集类 MyDataset
   # 2.数据集的操作
      # 使用 __len__ 方法获取数据集的长度，使用 __getitem__ 方法按索引获取单个样本
   # 3.数据集的使用
      # 创建数据集实例后，可以通过遍历的方式访问数据集中的每个样本，打印样本内容，便于检查数据格式和内容